import re

# todo: 暂未处理数据字典，因为表头格式不统一，需要和产品协调
def get_headings_and_formulas(ocr_response):
    res = []
    current_formula = []
    current_dict = []
    collecting_formula = False
    collecting_dict = False

    heading_pattern = re.compile(r'^#{1,4}\s(\d+(?:\.\d+)*)\s+(.+)$')
    formula_start_pattern = re.compile(r'【公式】')
    formula_end_pattern = re.compile(r'【说明】|^#{1,4}\s\d')
    dict_start_pattern = re.compile(r'^#{1,4}\s+(\d+(?:\.\d+)*)\s+数据字典$')
    dict_end_pattern = re.compile(r'^#{1,4}\s\d')

    for page in ocr_response.pages:
        lines = page.markdown.split('\n')
        for line in lines:
            line = line.strip()

            if match := heading_pattern.match(line):
                level, heading_text = match.groups()
                res.append({
                    "level": level,
                    "heading": heading_text,
                    "has_formula": False,
                    "formula": ""
                })

            if collecting_formula and formula_end_pattern.search(line):
                collecting_formula = False
                if res and current_formula:
                    res[-1]["formula"] = "\n".join(current_formula)
                    res[-1]["has_formula"] = True
                current_formula = []

            elif formula_start_pattern.search(line):
                collecting_formula = True

            elif collecting_formula:
                current_formula.append(line)

            if collecting_dict and dict_end_pattern.search(line):
                collecting_dict = False
            elif dict_start_pattern.search(line):
                collecting_dict = True
            elif collecting_dict:
                current_dict.append(line)

    current_dict = collect_dict(current_dict)

    return {"formulas": res, "dict": current_dict}

def collect_dict(current_dict):
    if len(current_dict) < 3:
        return []
    # 处理数据字典
    # todo: 数据字典表头可能不固定，需要根据表头名称“符号	维度	类型”来提取
    current_dict = [line.strip() for line in current_dict if line.strip()]
    current_dict = [line for line in current_dict if line.startswith('|')]
    dict_header = current_dict[0].split('|')
    dict_header = [h.strip() for h in dict_header if h.strip()]
    if "符号" not in dict_header or "维度" not in dict_header or "类型" not in dict_header:
        return []
    header_index = {h: i for i, h in enumerate(dict_header)}
    current_dict = current_dict[2:]
    dict_data = []
    for line in current_dict:
        parts = [part.strip() for part in line.split('|') if part.strip()]
        if len(parts) == 3:
            dict_data.append({
                "symbol": parts[header_index["符号"]],
                "dimension": parts[header_index["维度"]],
                "type": parts[header_index["类型"]]
            })

    return dict_data

import os
if __name__ == "__main__":
    home_dir = os.path.expanduser("~")
    pdf_path = os.path.join(home_dir, "Downloads", "文档导入-测试输入", "SOC文档.pdf")
    docx_path = os.path.join(home_dir, "Downloads", "文档导入-测试输入", "DSP文档.docx")
    pdf_path2 = os.path.join(home_dir, "Downloads", "文档导入-测试输入", "DSP文档_.pdf")

    # responses = call_ocr(pdf_path2)
    # # print(responses)
    # res = get_headings_and_formulas(responses)
    # for r in res:
    #     print(r)